{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "f1b23ffc-22d8-4414-b2e4-66d45a03523d",
   "metadata": {},
   "source": [
    "# LLM Reranker Demonstration (Great Gatsby)\n",
    "\n",
    "This tutorial showcases how to do a two-stage pass for retrieval. Use embedding-based retrieval with a high top-k value\n",
    "in order to maximize recall and get a large set of candidate items. Then, use LLM-based retrieval\n",
    "to dynamically select the nodes that are actually relevant to the query."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1bd392e1",
   "metadata": {},
   "outputs": [],
   "source": [
    "%pip install llama-index-llms-openai"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ab68a3ee-f4c1-4830-856a-9c37f7a42364",
   "metadata": {},
   "outputs": [],
   "source": [
    "import nest_asyncio\n",
    "\n",
    "nest_asyncio.apply()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8c16404b-db81-4ee6-9c0f-cc4cdf3cb147",
   "metadata": {},
   "outputs": [],
   "source": [
    "import logging\n",
    "import sys\n",
    "\n",
    "logging.basicConfig(stream=sys.stdout, level=logging.INFO)\n",
    "logging.getLogger().addHandler(logging.StreamHandler(stream=sys.stdout))\n",
    "from llama_index.core import VectorStoreIndex, SimpleDirectoryReader\n",
    "from llama_index.core.postprocessor import LLMRerank\n",
    "from llama_index.llms.openai import OpenAI\n",
    "from IPython.display import Markdown, display"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "8a5b0ac6-ec9c-40e1-9120-d20e33c37f80",
   "metadata": {},
   "source": [
    "## Load Data, Build Index"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b785b75a-cf4d-4799-8fd2-970e5c39727c",
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.core import Settings\n",
    "\n",
    "# LLM (gpt-3.5-turbo)\n",
    "Settings.llm = OpenAI(temperature=0, model=\"gpt-3.5-turbo\")\n",
    "Settings.chunk_size = 512"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9b689e92-fc55-48b5-859f-2a2c1394c872",
   "metadata": {},
   "outputs": [],
   "source": [
    "# load documents\n",
    "documents = SimpleDirectoryReader(\"../../../examples/gatsby/data\").load_data()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1a01180d-67c8-4664-a4ee-ceab1367cb78",
   "metadata": {},
   "outputs": [],
   "source": [
    "documents"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d8555b6d-5289-4803-bced-6b2f85dea3da",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "INFO:llama_index.token_counter.token_counter:> [build_index_from_nodes] Total LLM token usage: 0 tokens\n",
      "> [build_index_from_nodes] Total LLM token usage: 0 tokens\n",
      "INFO:llama_index.token_counter.token_counter:> [build_index_from_nodes] Total embedding token usage: 49266 tokens\n",
      "> [build_index_from_nodes] Total embedding token usage: 49266 tokens\n"
     ]
    }
   ],
   "source": [
    "index = VectorStoreIndex.from_documents(\n",
    "    documents,\n",
    ")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "805847d9-a2c1-4930-98a9-98126e730000",
   "metadata": {},
   "source": [
    "## Retrieval"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ace86835-832e-4964-a025-428891aa2c8b",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/var/folders/1r/c3h91d9s49xblwfvz79s78_c0000gn/T/ipykernel_44297/3519340226.py:7: FutureWarning: Passing a negative integer is deprecated in version 1.0 and will not be supported in future version. Instead, use None to not limit the column width.\n",
      "  pd.set_option('display.max_colwidth', -1)\n"
     ]
    }
   ],
   "source": [
    "from llama_index.core.retrievers import VectorIndexRetriever\n",
    "from llama_index.core import QueryBundle\n",
    "import pandas as pd\n",
    "from IPython.display import display, HTML\n",
    "\n",
    "\n",
    "pd.set_option(\"display.max_colwidth\", -1)\n",
    "\n",
    "\n",
    "def get_retrieved_nodes(\n",
    "    query_str, vector_top_k=10, reranker_top_n=3, with_reranker=False\n",
    "):\n",
    "    query_bundle = QueryBundle(query_str)\n",
    "    # configure retriever\n",
    "    retriever = VectorIndexRetriever(\n",
    "        index=index,\n",
    "        similarity_top_k=vector_top_k,\n",
    "    )\n",
    "    retrieved_nodes = retriever.retrieve(query_bundle)\n",
    "\n",
    "    if with_reranker:\n",
    "        # configure reranker\n",
    "        reranker = LLMRerank(\n",
    "            choice_batch_size=5,\n",
    "            top_n=reranker_top_n,\n",
    "        )\n",
    "        retrieved_nodes = reranker.postprocess_nodes(\n",
    "            retrieved_nodes, query_bundle\n",
    "        )\n",
    "\n",
    "    return retrieved_nodes\n",
    "\n",
    "\n",
    "def pretty_print(df):\n",
    "    return display(HTML(df.to_html().replace(\"\\\\n\", \"<br>\")))\n",
    "\n",
    "\n",
    "def visualize_retrieved_nodes(nodes) -> None:\n",
    "    result_dicts = []\n",
    "    for node in nodes:\n",
    "        result_dict = {\"Score\": node.score, \"Text\": node.node.get_text()}\n",
    "        result_dicts.append(result_dict)\n",
    "\n",
    "    pretty_print(pd.DataFrame(result_dicts))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "20993200-b725-403e-8a79-e2571dab2ebc",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "INFO:llama_index.token_counter.token_counter:> [retrieve] Total LLM token usage: 0 tokens\n",
      "> [retrieve] Total LLM token usage: 0 tokens\n",
      "INFO:llama_index.token_counter.token_counter:> [retrieve] Total embedding token usage: 10 tokens\n",
      "> [retrieve] Total embedding token usage: 10 tokens\n"
     ]
    }
   ],
   "source": [
    "new_nodes = get_retrieved_nodes(\n",
    "    \"Who was driving the car that hit Myrtle?\",\n",
    "    vector_top_k=3,\n",
    "    with_reranker=False,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9dd3bc25-6ef1-46b7-869e-247c83af9d4d",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>Score</th>\n",
       "      <th>Text</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>0.828844</td>\n",
       "      <td>and some garrulous man telling over and over what<br>had happened, until it became less and less real even to him and he<br>could tell it no longer, and Myrtle Wilson’s tragic achievement was<br>forgotten. Now I want to go back a little and tell what happened at<br>the garage after we left there the night before.<br><br>They had difficulty in locating the sister, Catherine. She must have<br>broken her rule against drinking that night, for when she arrived she<br>was stupid with liquor and unable to understand that the ambulance had<br>already gone to Flushing. When they convinced her of this, she<br>immediately fainted, as if that was the intolerable part of the<br>affair. Someone, kind or curious, took her in his car and drove her in<br>the wake of her sister’s body.<br><br>Until long after midnight a changing crowd lapped up against the front<br>of the garage, while George Wilson rocked himself back and forth on<br>the couch inside. For a while the door of the office was open, and<br>everyone who came into the garage glanced irresistibly through it.<br>Finally someone said it was a shame, and closed the door. Michaelis<br>and several other men were with him; first, four or five men, later<br>two or three men. Still later Michaelis had to ask the last stranger<br>to wait there fifteen minutes longer, while he went back to his own<br>place and made a pot of coffee. After that, he stayed there alone with<br>Wilson until dawn.<br><br>About three o’clock the quality of Wilson’s incoherent muttering<br>changed—he grew quieter and began to talk about the yellow car. He<br>announced that he had a way of finding out whom the yellow car<br>belonged to, and then he blurted out that a couple of months ago his<br>wife had come from the city with her face bruised and her nose<br>swollen.<br><br>But when he heard himself say this, he flinched and began to cry “Oh,<br>my God!” again in his groaning voice. Michaelis made a clumsy</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>0.827754</td>\n",
       "      <td>she rushed out into the dusk, waving her hands and<br>shouting—before he could move from his door the business was over.<br><br>The “death car” as the newspapers called it, didn’t stop; it came out<br>of the gathering darkness, wavered tragically for a moment, and then<br>disappeared around the next bend. Mavro Michaelis wasn’t even sure of<br>its colour—he told the first policeman that it was light green. The<br>other car, the one going toward New York, came to rest a hundred yards<br>beyond, and its driver hurried back to where Myrtle Wilson, her life<br>violently extinguished, knelt in the road and mingled her thick dark<br>blood with the dust.<br><br>Michaelis and this man reached her first, but when they had torn open<br>her shirtwaist, still damp with perspiration, they saw that her left<br>breast was swinging loose like a flap, and there was no need to listen<br>for the heart beneath. The mouth was wide open and ripped a little at<br>the corners, as though she had choked a little in giving up the<br>tremendous vitality she had stored so long.<br><br>------------------------------------------------------------------------<br><br>We saw the three or four automobiles and the crowd when we were still<br>some distance away.<br><br>“Wreck!” said Tom. “That’s good. Wilson’ll have a little business at<br>last.”<br><br>He slowed down, but still without any intention of stopping, until, as<br>we came nearer, the hushed, intent faces of the people at the garage<br>door made him automatically put on the brakes.<br><br>“We’ll take a look,” he said doubtfully, “just a look.”<br><br>I became aware now of a hollow, wailing sound which issued incessantly<br>from the garage, a sound which as we got out of the coupé and walked<br>toward the door resolved itself into the words “Oh, my God!” uttered<br>over and over in a gasping</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>0.826390</td>\n",
       "      <td>went on, “and left the car in<br>my garage. I don’t think anybody saw us, but of course I can’t be<br>sure.”<br><br>I disliked him so much by this time that I didn’t find it necessary to<br>tell him he was wrong.<br><br>“Who was the woman?” he inquired.<br><br>“Her name was Wilson. Her husband owns the garage. How the devil did<br>it happen?”<br><br>“Well, I tried to swing the wheel—” He broke off, and suddenly I<br>guessed at the truth.<br><br>“Was Daisy driving?”<br><br>“Yes,” he said after a moment, “but of course I’ll say I was. You see,<br>when we left New York she was very nervous and she thought it would<br>steady her to drive—and this woman rushed out at us just as we were<br>passing a car coming the other way. It all happened in a minute, but<br>it seemed to me that she wanted to speak to us, thought we were<br>somebody she knew. Well, first Daisy turned away from the woman toward<br>the other car, and then she lost her nerve and turned back. The second<br>my hand reached the wheel I felt the shock—it must have killed her<br>instantly.”<br><br>“It ripped her open—”<br><br>“Don’t tell me, old sport.” He winced. “Anyhow—Daisy stepped on it. I<br>tried to make her stop, but she couldn’t, so I pulled on the emergency<br>brake. Then she fell over into my lap and I drove on.<br><br>“She’ll be all right tomorrow,” he said presently. “I’m just going to<br>wait here and see if he tries to bother her about that unpleasantness<br>this afternoon. She’s locked herself into her room, and if he tries<br>any brutality she’s going to turn the light out and on again.”<br><br>“He won’t touch</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "visualize_retrieved_nodes(new_nodes)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0ad7d62b-dba9-45d5-896f-4baa9a40edba",
   "metadata": {},
   "outputs": [],
   "source": [
    "new_nodes = get_retrieved_nodes(\n",
    "    \"Who was driving the car that hit Myrtle?\",\n",
    "    vector_top_k=10,\n",
    "    reranker_top_n=3,\n",
    "    with_reranker=True,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ae22d297-dd3b-4a71-a890-faceafab4e1a",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>Score</th>\n",
       "      <th>Text</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>10.0</td>\n",
       "      <td>went on, “and left the car in<br>my garage. I don’t think anybody saw us, but of course I can’t be<br>sure.”<br><br>I disliked him so much by this time that I didn’t find it necessary to<br>tell him he was wrong.<br><br>“Who was the woman?” he inquired.<br><br>“Her name was Wilson. Her husband owns the garage. How the devil did<br>it happen?”<br><br>“Well, I tried to swing the wheel—” He broke off, and suddenly I<br>guessed at the truth.<br><br>“Was Daisy driving?”<br><br>“Yes,” he said after a moment, “but of course I’ll say I was. You see,<br>when we left New York she was very nervous and she thought it would<br>steady her to drive—and this woman rushed out at us just as we were<br>passing a car coming the other way. It all happened in a minute, but<br>it seemed to me that she wanted to speak to us, thought we were<br>somebody she knew. Well, first Daisy turned away from the woman toward<br>the other car, and then she lost her nerve and turned back. The second<br>my hand reached the wheel I felt the shock—it must have killed her<br>instantly.”<br><br>“It ripped her open—”<br><br>“Don’t tell me, old sport.” He winced. “Anyhow—Daisy stepped on it. I<br>tried to make her stop, but she couldn’t, so I pulled on the emergency<br>brake. Then she fell over into my lap and I drove on.<br><br>“She’ll be all right tomorrow,” he said presently. “I’m just going to<br>wait here and see if he tries to bother her about that unpleasantness<br>this afternoon. She’s locked herself into her room, and if he tries<br>any brutality she’s going to turn the light out and on again.”<br><br>“He won’t touch</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "visualize_retrieved_nodes(new_nodes)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4e540fb6-189f-4a54-98d6-852e561f39ee",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "INFO:llama_index.token_counter.token_counter:> [retrieve] Total LLM token usage: 0 tokens\n",
      "> [retrieve] Total LLM token usage: 0 tokens\n",
      "INFO:llama_index.token_counter.token_counter:> [retrieve] Total embedding token usage: 14 tokens\n",
      "> [retrieve] Total embedding token usage: 14 tokens\n"
     ]
    }
   ],
   "source": [
    "new_nodes = get_retrieved_nodes(\n",
    "    \"What did Gatsby want Daisy to do in front of Tom?\",\n",
    "    vector_top_k=3,\n",
    "    with_reranker=False,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3b82380c-e29f-4e9e-a46d-038ca1469904",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "\n",
      "****Score****: 0.8647796939111776\n",
      "****Node text****\n",
      ": got to make your house into a pigsty in order to have any\n",
      "friends—in the modern world.”\n",
      "\n",
      "Angry as I was, as we all were, I was tempted to laugh whenever he\n",
      "opened his mouth. The transition from libertine to prig was so\n",
      "complete.\n",
      "\n",
      "“I’ve got something to tell you, old sport—” began Gatsby. But Daisy\n",
      "guessed at his intention.\n",
      "\n",
      "“Please don’t!” she interrupted helplessly. “Please let’s all go\n",
      "home. Why don’t we all go home?”\n",
      "\n",
      "“That’s a good idea,” I got up. “Come on, Tom. Nobody wants a drink.”\n",
      "\n",
      "“I want to know what Mr. Gatsby has to tell me.”\n",
      "\n",
      "“Your wife doesn’t love you,” said Gatsby. “She’s never loved you.\n",
      "She loves me.”\n",
      "\n",
      "“You must be crazy!” exclaimed Tom automatically.\n",
      "\n",
      "Gatsby sprang to his feet, vivid with excitement.\n",
      "\n",
      "“She never loved you, do you hear?” he cried. “She only married you\n",
      "because I was poor and she was tired of waiting for me. It was a\n",
      "terrible mistake, but in her heart she never loved anyone except me!”\n",
      "\n",
      "At this point Jordan and I tried to go, but Tom and Gatsby insisted\n",
      "with competitive firmness that we remain—as though neither of them had\n",
      "anything to conceal and it would be a privilege to partake vicariously\n",
      "of their emotions.\n",
      "\n",
      "“Sit down, Daisy,” Tom’s voice groped unsuccessfully for the paternal\n",
      "note. “What’s been going on? I want to hear all about it.”\n",
      "\n",
      "“I told you what’s been going on,” said Gatsby. “Going on for five\n",
      "years—and you didn’t know.”\n",
      "\n",
      "Tom turned to Daisy\n",
      "\n",
      "\n",
      "****Score****: 0.8609230717744326\n",
      "****Node text****\n",
      ": to keep your\n",
      "shoes dry?” There was a husky tenderness in his tone … “Daisy?”\n",
      "\n",
      "“Please don’t.” Her voice was cold, but the rancour was gone from it.\n",
      "She looked at Gatsby. “There, Jay,” she said—but her hand as she tried\n",
      "to light a cigarette was trembling. Suddenly she threw the cigarette\n",
      "and the burning match on the carpet.\n",
      "\n",
      "“Oh, you want too much!” she cried to Gatsby. “I love you now—isn’t\n",
      "that enough? I can’t help what’s past.” She began to sob\n",
      "helplessly. “I did love him once—but I loved you too.”\n",
      "\n",
      "Gatsby’s eyes opened and closed.\n",
      "\n",
      "“You loved me too?” he repeated.\n",
      "\n",
      "“Even that’s a lie,” said Tom savagely. “She didn’t know you were\n",
      "alive. Why—there’s things between Daisy and me that you’ll never know,\n",
      "things that neither of us can ever forget.”\n",
      "\n",
      "The words seemed to bite physically into Gatsby.\n",
      "\n",
      "“I want to speak to Daisy alone,” he insisted. “She’s all excited\n",
      "now—”\n",
      "\n",
      "“Even alone I can’t say I never loved Tom,” she admitted in a pitiful\n",
      "voice. “It wouldn’t be true.”\n",
      "\n",
      "“Of course it wouldn’t,” agreed Tom.\n",
      "\n",
      "She turned to her husband.\n",
      "\n",
      "“As if it mattered to you,” she said.\n",
      "\n",
      "“Of course it matters. I’m going to take better care of you from now\n",
      "on.”\n",
      "\n",
      "“You don’t understand,” said Gatsby, with a touch of panic. “You’re\n",
      "not going to take care of her any more.”\n",
      "\n",
      "“I’m not?” Tom opened his eyes wide and\n",
      "\n",
      "\n",
      "****Score****: 0.8555028907426916\n",
      "****Node text****\n",
      ": shadowed well with awnings, was dark and cool. Daisy and\n",
      "Jordan lay upon an enormous couch, like silver idols weighing down\n",
      "their own white dresses against the singing breeze of the fans.\n",
      "\n",
      "“We can’t move,” they said together.\n",
      "\n",
      "Jordan’s fingers, powdered white over their tan, rested for a moment\n",
      "in mine.\n",
      "\n",
      "“And Mr. Thomas Buchanan, the athlete?” I inquired.\n",
      "\n",
      "Simultaneously I heard his voice, gruff, muffled, husky, at the hall\n",
      "telephone.\n",
      "\n",
      "Gatsby stood in the centre of the crimson carpet and gazed around with\n",
      "fascinated eyes. Daisy watched him and laughed, her sweet, exciting\n",
      "laugh; a tiny gust of powder rose from her bosom into the air.\n",
      "\n",
      "“The rumour is,” whispered Jordan, “that that’s Tom’s girl on the\n",
      "telephone.”\n",
      "\n",
      "We were silent. The voice in the hall rose high with annoyance: “Very\n",
      "well, then, I won’t sell you the car at all … I’m under no obligations\n",
      "to you at all … and as for your bothering me about it at lunch time, I\n",
      "won’t stand that at all!”\n",
      "\n",
      "“Holding down the receiver,” said Daisy cynically.\n",
      "\n",
      "“No, he’s not,” I assured her. “It’s a bona-fide deal. I happen to\n",
      "know about it.”\n",
      "\n",
      "Tom flung open the door, blocked out its space for a moment with his\n",
      "thick body, and hurried into the room.\n",
      "\n",
      "“Mr. Gatsby!” He put out his broad, flat hand with well-concealed\n",
      "dislike. “I’m glad to see you, sir … Nick …”\n",
      "\n",
      "“Make us a cold drink,” cried Daisy.\n",
      "\n",
      "As he left the room again she got up and went over to Gatsby and\n",
      "pulled his face\n"
     ]
    }
   ],
   "source": [
    "visualize_retrieved_nodes(new_nodes)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "aec75734-52e4-4b7e-9b6d-d0b7306ae5d8",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "INFO:llama_index.token_counter.token_counter:> [retrieve] Total LLM token usage: 0 tokens\n",
      "> [retrieve] Total LLM token usage: 0 tokens\n",
      "INFO:llama_index.token_counter.token_counter:> [retrieve] Total embedding token usage: 14 tokens\n",
      "> [retrieve] Total embedding token usage: 14 tokens\n",
      "Doc: 2, Relevance: 10\n",
      "No relevant documents found. Please provide a different question.\n"
     ]
    }
   ],
   "source": [
    "new_nodes = get_retrieved_nodes(\n",
    "    \"What did Gatsby want Daisy to do in front of Tom?\",\n",
    "    vector_top_k=10,\n",
    "    reranker_top_n=3,\n",
    "    with_reranker=True,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "89e5408e-bf4b-483f-9144-fc5ac4fb48e8",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "\n",
      "****Score****: 10.0\n",
      "****Node text****\n",
      ": to keep your\n",
      "shoes dry?” There was a husky tenderness in his tone … “Daisy?”\n",
      "\n",
      "“Please don’t.” Her voice was cold, but the rancour was gone from it.\n",
      "She looked at Gatsby. “There, Jay,” she said—but her hand as she tried\n",
      "to light a cigarette was trembling. Suddenly she threw the cigarette\n",
      "and the burning match on the carpet.\n",
      "\n",
      "“Oh, you want too much!” she cried to Gatsby. “I love you now—isn’t\n",
      "that enough? I can’t help what’s past.” She began to sob\n",
      "helplessly. “I did love him once—but I loved you too.”\n",
      "\n",
      "Gatsby’s eyes opened and closed.\n",
      "\n",
      "“You loved me too?” he repeated.\n",
      "\n",
      "“Even that’s a lie,” said Tom savagely. “She didn’t know you were\n",
      "alive. Why—there’s things between Daisy and me that you’ll never know,\n",
      "things that neither of us can ever forget.”\n",
      "\n",
      "The words seemed to bite physically into Gatsby.\n",
      "\n",
      "“I want to speak to Daisy alone,” he insisted. “She’s all excited\n",
      "now—”\n",
      "\n",
      "“Even alone I can’t say I never loved Tom,” she admitted in a pitiful\n",
      "voice. “It wouldn’t be true.”\n",
      "\n",
      "“Of course it wouldn’t,” agreed Tom.\n",
      "\n",
      "She turned to her husband.\n",
      "\n",
      "“As if it mattered to you,” she said.\n",
      "\n",
      "“Of course it matters. I’m going to take better care of you from now\n",
      "on.”\n",
      "\n",
      "“You don’t understand,” said Gatsby, with a touch of panic. “You’re\n",
      "not going to take care of her any more.”\n",
      "\n",
      "“I’m not?” Tom opened his eyes wide and\n"
     ]
    }
   ],
   "source": [
    "visualize_retrieved_nodes(new_nodes)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "23dcdac6-f4dd-469e-9f47-d030f27bacda",
   "metadata": {},
   "source": [
    "## Query Engine"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "94f6531d-953a-41c6-9bcc-9c293550f8c4",
   "metadata": {},
   "outputs": [],
   "source": [
    "query_engine = index.as_query_engine(\n",
    "    similarity_top_k=10,\n",
    "    node_postprocessors=[\n",
    "        LLMRerank(\n",
    "            choice_batch_size=5,\n",
    "            top_n=2,\n",
    "        )\n",
    "    ],\n",
    "    response_mode=\"tree_summarize\",\n",
    ")\n",
    "response = query_engine.query(\n",
    "    \"What did the author do during his time at Y Combinator?\",\n",
    ")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "llama_index_v2",
   "language": "python",
   "name": "llama_index_v2"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
